close all; clear; clc;

Alpha = [0, pi/2, 0, pi/2, -pi/2, pi/2];
A = [0, 0.05, 0.50, 0.05, 0, 0];
D = [0.50, 0, 0, 0.50, 0, 0.1];
Theta = [0, pi/2, 0, 0, 0, 0];

for i = 1:6
    L(i) = Link('alpha', Alpha(i),'a', A(i),'d', D(i)); L(i).offset = Theta(i); L(i).qlim=[-2 * pi, 2 * pi]; L(i).mdh = 1;
end

L(1).m = 1; L(1).r = [A(1), 0, D(1)]; L(1).I = eye(3);
L(2).m = 1; L(2).r = [A(2), 0, D(2)]; L(2).I = eye(3);
L(3).m = 1; L(3).r = [A(3), 0, D(3)]; L(3).I = eye(3);
L(4).m = 0.5; L(4).r = [A(4), 0, D(4)]; L(4).I = 0.05 * eye(3);
L(5).m = 0.1; L(5).r = [A(5), 0, D(5)]; L(5).I = 0.01 * eye(3);
L(6).m = 0.1; L(6).r = [A(6), 0, D(6)]; L(6).I = 0.01 * eye(3);

robot = SerialLink(L, 'name', 'robot');
robot.gravity = [0; 0; 0];

q1_0 = 7.65404249470431e-18;
q2_0 = -0.520469582664616;
q3_0 = 0.353499793119520;
q4_0 = 1.24191828169582e-16;
q5_0 = -1.40382653724961;
q6_0 = -1.29860243047315e-17;
q0 = [q1_0, q2_0, q3_0, q4_0, q5_0, q6_0];

dt = 0.1;
tf = 5;
t = [0:dt:tf]';

len = length(t);
k = zeros(len, 1);
T = zeros(len, 6);

M = 5;
times = zeros(M+1, 1);
emax = zeros(M+1, 6);
emean = zeros(M+1, 6);

for i = 0:1:M
    pause(0.01);
    
    simOut = sim('sim_main', [0, tf]);
    
    t = simOut.t.data;
    T = simOut.T.data;
    e = simOut.e.data;
    q_d = simOut.q_d.data;
    dq_d = simOut.dq_d.data;
    q = simOut.q.data;
    dq = simOut.dq.data;
    px = simOut.px.data;
    py = simOut.py.data;
    pz = simOut.pz.data;
    
    figure(1);
    for j = 1:6
        subplot(3, 2, j);
        hold on;
        plot(t, q_d(:, j), 'r', t, q(:, j), 'b--');
        xlabel("time (s)");
        ylabel("q" + j + " (rad)");
    end
    
    if i == 0
        q_0 = q;
    end
    
    j = i + 1;
    times(j) = i;
    emax(j, :) = max(abs(e));
    emean(j, :) = mean(abs(e));
end

figure(2);
for j = 1:2
    subplot(2, 1, j);
    hold on;
    plot(t, q_d(:, j + 3), 'r', t, q_0(:, j + 3), 'g-.', t, q(:, j + 3), 'b--');
    legend("desired", "before learning", "after learning");
    xlabel("time (s)");
    ylabel("q" + (j + 3) + " (rad)");
end

figure(3);
for j = 1:2
    subplot(2, 1, j);
    hold on;
    plot(times, emax(:, j + 3), '* - b');
    xlabel("times");
    ylabel("q" + (j + 3) + " error");
end

figure(4)
plot3(px, py, pz);
robot.plot(q);